//===--- CapturePropagation.cpp - Propagate closure capture constants -----===//
//
// This source file is part of the Swift.org open source project
//
// Copyright (c) 2014 - 2017 Apple Inc. and the Swift project authors
// Licensed under Apache License v2.0 with Runtime Library Exception
//
// See https://swift.org/LICENSE.txt for license information
// See https://swift.org/CONTRIBUTORS.txt for the list of Swift project authors
//
//===----------------------------------------------------------------------===//

#define DEBUG_TYPE "capture-prop"

#include "polarphp/ast/GenericEnvironment.h"
#include "polarphp/demangling/Demangler.h"
#include "polarphp/pil/lang/PILCloner.h"
#include "polarphp/pil/lang/PILInstruction.h"
#include "polarphp/pil/lang/TypeSubstCloner.h"
#include "polarphp/pil/optimizer/analysis/ColdBlockInfo.h"
#include "polarphp/pil/optimizer/analysis/DominanceAnalysis.h"
#include "polarphp/pil/optimizer/passmgr/Transforms.h"
#include "polarphp/pil/optimizer/utils/Generics.h"
#include "polarphp/pil/optimizer/utils/InstOptUtils.h"
#include "polarphp/pil/optimizer/utils/PILOptFunctionBuilder.h"
#include "polarphp/pil/optimizer/utils/SpecializationMangler.h"
#include "llvm/ADT/MapVector.h"
#include "llvm/ADT/Statistic.h"
#include "llvm/Support/Debug.h"

using namespace polar;

STATISTIC(NumCapturesPropagated, "Number of constant captures propagated");

namespace {
/// Propagate constants through closure captures by specializing the partially
/// applied function.
/// Also optimize away partial_apply instructions where all partially applied
/// arguments are dead.
class CapturePropagation : public PILFunctionTransform
{
public:
   void run() override;

protected:
   bool optimizePartialApply(PartialApplyInst *PAI);
   PILFunction *specializeConstClosure(PartialApplyInst *PAI,
                                       PILFunction *SubstF);
   void rewritePartialApply(PartialApplyInst *PAI, PILFunction *SpecialF);
};
} // end anonymous namespace

static LiteralInst *getConstant(PILValue V) {
   if (auto I = dyn_cast<ThinToThickFunctionInst>(V))
      return getConstant(I->getOperand());
   if (auto I = dyn_cast<ConvertFunctionInst>(V))
      return getConstant(I->getOperand());
   return dyn_cast<LiteralInst>(V);
}

static bool isOptimizableConstant(PILValue V) {
   // We do not optimize string literals of length > 32 since we would need to
   // encode them into the symbol name for uniqueness.
   if (auto *SLI = dyn_cast<StringLiteralInst>(V))
      return SLI->getValue().size() <= 32;
   return true;
}

static bool isConstant(PILValue V) {
   V = getConstant(V);
   return V && isOptimizableConstant(V);
}

static std::string getClonedName(PartialApplyInst *PAI, IsSerialized_t Serialized,
                                 PILFunction *F) {
   auto P = demangling::SpecializationPass::CapturePropagation;
   mangle::FunctionSignatureSpecializationMangler Mangler(P, Serialized, F);

   // We know that all arguments are literal insts.
   unsigned argIdx = ApplySite(PAI).getCalleeArgIndexOfFirstAppliedArg();
   for (auto arg : PAI->getArguments()) {
      Mangler.setArgumentConstantProp(argIdx, getConstant(arg));
      ++argIdx;
   }
   return Mangler.mangle();
}

namespace {
/// Clone the partially applied function, replacing incoming arguments with
/// literal constants.
///
/// The cloned literals will retain the PILLocation from the partial apply's
/// caller, so the cloned function will have a mix of locations from different
/// functions.
class CapturePropagationCloner
   : public TypeSubstCloner<CapturePropagationCloner, PILOptFunctionBuilder> {
   using SuperTy =
   TypeSubstCloner<CapturePropagationCloner, PILOptFunctionBuilder>;
   friend class PILInstructionVisitor<CapturePropagationCloner>;
   friend class PILCloner<CapturePropagationCloner>;

   PILFunction *OrigF;
   bool IsCloningConstant;
public:
   CapturePropagationCloner(PILFunction *OrigF, PILFunction *NewF,
                            SubstitutionMap Subs)
      : SuperTy(*NewF, *OrigF, Subs), OrigF(OrigF), IsCloningConstant(false) {}

   void cloneClosure(OperandValueArrayRef Args);

protected:
   /// Literals cloned from the caller drop their location so the debug line
   /// tables don't senselessly jump around. As a placeholder give them the
   /// location of the newly cloned function.
   PILLocation remapLocation(PILLocation InLoc) {
      if (IsCloningConstant)
         return getBuilder().getFunction().getLocation();
      return InLoc;
   }

   /// Literals cloned from the caller take on the new function's debug scope.
   void postProcess(PILInstruction *Orig, PILInstruction *Cloned) {
      assert(IsCloningConstant == (Orig->getFunction() != OrigF) &&
             "Expect only cloned constants from the caller function.");
      PILClonerWithScopes<CapturePropagationCloner>::postProcess(Orig, Cloned);
   }

   const PILDebugScope *remapScope(const PILDebugScope *DS) {
      if (IsCloningConstant)
         return getBuilder().getFunction().getDebugScope();
      else
         return PILClonerWithScopes<CapturePropagationCloner>::remapScope(DS);
   }

   void cloneConstValue(PILValue Const);
};
} // end anonymous namespace

/// Clone a constant value. Recursively walk the operand chain through cast
/// instructions to ensure that all dependents are cloned. Note that the
/// original value may not belong to the same function as the one being cloned
/// by cloneClosure() (they may be from the partial apply caller).
void CapturePropagationCloner::cloneConstValue(PILValue Val) {
   assert(IsCloningConstant && "incorrect mode");

   if (isValueCloned(Val))
      return;

   // TODO: MultiValueInstruction?
   auto Inst = dyn_cast<SingleValueInstruction>(Val);
   if (!Inst)
      return;

   if (Inst->getNumOperands() > 0) {
      // Only handle single operands for simple recursion without a worklist.
      assert(Inst->getNumOperands() == 1 && "expected single-operand cast");
      cloneConstValue(Inst->getOperand(0));
   }
   visit(Inst);
}

/// Clone the original partially applied function into the new specialized
/// function, replacing some arguments with literals.
void CapturePropagationCloner::cloneClosure(
   OperandValueArrayRef PartialApplyArgs) {

   PILFunction &CloneF = getBuilder().getFunction();

   // Create the entry basic block with the function arguments.
   PILBasicBlock *OrigEntryBB = &*OrigF->begin();
   PILBasicBlock *ClonedEntryBB = CloneF.createBasicBlock();
   auto cloneConv = CloneF.getConventions();

   // Only clone the arguments that remain in the new function type. The trailing
   // arguments are now propagated through the partial apply.
   assert(!IsCloningConstant && "incorrect mode");

   SmallVector<PILValue, 4> entryArgs;
   entryArgs.reserve(OrigEntryBB->getArguments().size());

   unsigned ArgIdx = 0;
   for (unsigned NewArgEnd = cloneConv.getNumPILArguments(); ArgIdx != NewArgEnd;
        ++ArgIdx) {

      PILArgument *Arg = OrigEntryBB->getArgument(ArgIdx);

      PILValue MappedValue = ClonedEntryBB->createFunctionArgument(
         remapType(Arg->getType()), Arg->getDecl());
      entryArgs.push_back(MappedValue);
   }
   assert(OrigEntryBB->args_size() - ArgIdx == PartialApplyArgs.size()
          && "unexpected number of partial apply arguments");

   // Replace the rest of the old arguments with constants.
   getBuilder().setInsertionPoint(ClonedEntryBB);
   IsCloningConstant = true;
   for (PILValue PartialApplyArg : PartialApplyArgs) {
      assert(isConstant(PartialApplyArg) &&
             "expected a constant arg to partial apply");

      cloneConstValue(PartialApplyArg);

      // The PartialApplyArg from the caller is now mapped to its cloned
      // instruction.  Also map the original argument to the cloned instruction.
      entryArgs.push_back(getMappedValue(PartialApplyArg));
      ++ArgIdx;
   }
   IsCloningConstant = false;

   // Clear information about cloned values from the caller function.
   clearClonerState();

   // Visit original BBs in depth-first preorder, starting with the
   // entry block, cloning all instructions and terminators.
   cloneFunctionBody(OrigF, ClonedEntryBB, entryArgs);
}

CanPILFunctionType getPartialApplyInterfaceResultType(PartialApplyInst *PAI) {
   // The new partial_apply will no longer take any arguments--they are all
   // expressed as literals. So its callee signature will be the same as its
   // return signature.
   auto FTy = PAI->getType().castTo<PILFunctionType>();
   assert(!PAI->hasSubstitutions() ||
          !PAI->getSubstitutionMap().hasArchetypes());
   FTy = cast<PILFunctionType>(
      FTy->mapTypeOutOfContext()->getCanonicalType());
   auto NewFTy = FTy;
   return NewFTy;
}

/// Given a partial_apply instruction, create a specialized callee by removing
/// all constant arguments and adding constant literals to the specialized
/// function body.
PILFunction *CapturePropagation::specializeConstClosure(PartialApplyInst *PAI,
                                                        PILFunction *OrigF) {
   IsSerialized_t Serialized = IsNotSerialized;
   if (PAI->getFunction()->isSerialized() && OrigF->isSerialized())
      Serialized = IsSerializable;

   std::string Name = getClonedName(PAI, Serialized, OrigF);

   // See if we already have a version of this function in the module. If so,
   // just return it.
   if (auto *NewF = OrigF->getModule().lookUpFunction(Name)) {
      assert(NewF->isSerialized() == Serialized);
      LLVM_DEBUG(llvm::dbgs()
                    << "  Found an already specialized version of the callee: ";
                    NewF->printName(llvm::dbgs()); llvm::dbgs() << "\n");
      return NewF;
   }

   // The new partial_apply will no longer take any arguments--they are all
   // expressed as literals. So its callee signature will be the same as its
   // return signature.
   auto NewFTy = getPartialApplyInterfaceResultType(PAI);
   NewFTy = NewFTy->getWithRepresentation(PILFunctionType::Representation::Thin);

   GenericEnvironment *GenericEnv = nullptr;
   if (NewFTy->getInvocationGenericSignature())
      GenericEnv = OrigF->getGenericEnvironment();
   PILOptFunctionBuilder FuncBuilder(*this);
   PILFunction *NewF = FuncBuilder.createFunction(
      PILLinkage::Shared, Name, NewFTy, GenericEnv, OrigF->getLocation(),
      OrigF->isBare(), OrigF->isTransparent(), Serialized, IsNotDynamic,
      OrigF->getEntryCount(), OrigF->isThunk(), OrigF->getClassSubclassScope(),
      OrigF->getInlineStrategy(), OrigF->getEffectsKind(),
      /*InsertBefore*/ OrigF, OrigF->getDebugScope());
   if (!OrigF->hasOwnership()) {
      NewF->setOwnershipEliminated();
   }
   LLVM_DEBUG(llvm::dbgs() << "  Specialize callee as ";
                 NewF->printName(llvm::dbgs());
                 llvm::dbgs() << " " << NewFTy << "\n");

   LLVM_DEBUG(if (PAI->hasSubstitutions()) {
      llvm::dbgs() << "CapturePropagation of generic partial_apply:\n";
      PAI->dumpInContext();
   });
   CapturePropagationCloner cloner(OrigF, NewF, PAI->getSubstitutionMap());
   cloner.cloneClosure(PAI->getArguments());
   assert(OrigF->getDebugScope()->Parent != NewF->getDebugScope()->Parent);
   return NewF;
}

void CapturePropagation::rewritePartialApply(PartialApplyInst *OrigPAI,
                                             PILFunction *SpecialF) {
   LLVM_DEBUG(llvm::dbgs() << "\n  Rewriting a partial apply:\n";
                 OrigPAI->dumpInContext();
                 llvm::dbgs() << "   with special function: "
                              << SpecialF->getName() << "\n";
                 llvm::dbgs() << "\nThe function being rewritten is:\n";
                 OrigPAI->getFunction()->dump());

   PILBuilderWithScope Builder(OrigPAI);
   auto FuncRef = Builder.createFunctionRef(OrigPAI->getLoc(), SpecialF);
   auto *T2TF = Builder.createThinToThickFunction(OrigPAI->getLoc(), FuncRef,
                                                  OrigPAI->getType());
   OrigPAI->replaceAllUsesWith(T2TF);
   // Remove any dealloc_stack users.
   SmallVector<Operand*, 16> Uses(T2TF->getUses());
   for (auto *Use : Uses)
      if (auto *DS = dyn_cast<DeallocStackInst>(Use->getUser()))
         DS->eraseFromParent();
   recursivelyDeleteTriviallyDeadInstructions(OrigPAI, true);
   LLVM_DEBUG(llvm::dbgs() << "  Rewrote caller:\n" << *T2TF);
}

/// For now, we conservative only specialize if doing so can eliminate dynamic
/// dispatch.
///
/// TODO: Check for other profitable constant propagation, like builtin compare.
static bool isProfitable(PILFunction *Callee) {
   PILBasicBlock *EntryBB = &*Callee->begin();
   for (auto *Arg : EntryBB->getArguments()) {
      for (auto *Operand : Arg->getUses()) {
         if (FullApplySite FAS = FullApplySite::isa(Operand->getUser())) {
            if (FAS.getCallee() == Operand->get())
               return true;
         }
      }
   }
   return false;
}

/// Returns true if block \p BB only contains a return or throw of the first
/// block argument and side-effect-free instructions.
static bool onlyContainsReturnOrThrowOfArg(PILBasicBlock *BB) {
   for (PILInstruction &I : *BB) {
      if (isa<ReturnInst>(&I) || isa<ThrowInst>(&I)) {
         PILValue RetVal = I.getOperand(0);
         return BB->getNumArguments() == 1 && RetVal == BB->getArgument(0);
      }
      if (I.mayHaveSideEffects() || isa<TermInst>(&I))
         return false;
   }
   llvm_unreachable("should have seen a terminator instruction");
}

/// Checks if \p Orig is a thunk which calls another function but without
/// passing the trailing \p numDeadParams dead parameters.
/// If a generic specialization was performed for a generic capture,
/// GenericSpecialized contains a tuple:
/// (new specialized function, old function)
static PILFunction *getSpecializedWithDeadParams(
   PILOptFunctionBuilder &FuncBuilder,
   PartialApplyInst *PAI, PILFunction *Orig, int numDeadParams,
   std::pair<PILFunction *, PILFunction *> &GenericSpecialized) {
   PILBasicBlock &EntryBB = *Orig->begin();
   unsigned NumArgs = EntryBB.getNumArguments();

   // Check if all dead parameters have trivial types. We don't support non-
   // trivial types because it's very hard to find places where we can release
   // those parameters (as a replacement for the removed partial_apply).
   // TODO: maybe we can skip this restriction when we have semantic ARC.
   for (unsigned Idx = NumArgs - numDeadParams; Idx < NumArgs; ++Idx) {
      PILType ArgTy = EntryBB.getArgument(Idx)->getType();
      if (!ArgTy.isTrivial(*Orig))
         return nullptr;
   }
   PILFunction *Specialized = nullptr;
   PILValue RetValue;

   // Check all instruction of the entry block.
   for (PILInstruction &I : EntryBB) {
      if (auto FAS = FullApplySite::isa(&I)) {
         // Check if this is the call of the specialized function.
         // If the original partial_apply didn't have substitutions,
         // also the specialized function must be not generic.
         if (!PAI->hasSubstitutions() && FAS.hasSubstitutions())
            return nullptr;

         // Is it the only call?
         if (Specialized)
            return nullptr;

         Specialized = FAS.getReferencedFunctionOrNull();
         if (!Specialized)
            return nullptr;

         // Check if parameters are passes 1-to-1
         unsigned NumArgs = FAS.getNumArguments();
         if (EntryBB.getNumArguments() - numDeadParams != NumArgs)
            return nullptr;

         for (unsigned Idx = 0; Idx < NumArgs; ++Idx) {
            if (FAS.getArgument(Idx) != (ValueBase *)EntryBB.getArgument(Idx))
               return nullptr;
         }

         if (auto *TAI = dyn_cast<TryApplyInst>(&I)) {
            // Check the normal and throw blocks of the try_apply.
            if (onlyContainsReturnOrThrowOfArg(TAI->getNormalBB()) &&
                onlyContainsReturnOrThrowOfArg(TAI->getErrorBB()))
               return Specialized;
            return nullptr;
         }
         assert(isa<ApplyInst>(&I) && "unknown FullApplySite instruction");
         RetValue = cast<ApplyInst>(&I);
         continue;
      }
      if (auto *RI = dyn_cast<ReturnInst>(&I)) {
         // Check if we return the result of the apply.
         if (RI->getOperand() != RetValue)
            return nullptr;
         continue;
      }
      if (I.mayHaveSideEffects() || isa<TermInst>(&I))
         return nullptr;
   }

   auto Rep = Specialized->getLoweredFunctionType()->getRepresentation();
   if (getPILFunctionLanguage(Rep) != PILFunctionLanguage::Swift)
      return nullptr;

   GenericSpecialized = std::make_pair(nullptr, nullptr);

   if (PAI->hasSubstitutions()) {
      if (Specialized->isExternalDeclaration())
         return nullptr;
      if (!Orig->shouldOptimize())
         return nullptr;

      // Perform a generic specialization of the Specialized function.
      ReabstractionInfo ReInfo(
         FuncBuilder.getModule().getTypePHPModule(),
         FuncBuilder.getModule().isWholeModule(), ApplySite(), Specialized,
         PAI->getSubstitutionMap(), Specialized->isSerialized(),
         /* ConvertIndirectToDirect */ false);
      GenericFuncSpecializer FuncSpecializer(FuncBuilder,
                                             Specialized,
                                             ReInfo.getClonerParamSubstitutionMap(),
                                             ReInfo);

      PILFunction *GenericSpecializedFunc = FuncSpecializer.trySpecialization();
      if (!GenericSpecializedFunc)
         return nullptr;
      GenericSpecialized = std::make_pair(GenericSpecializedFunc, Specialized);
      return GenericSpecializedFunc;
   }
   return Specialized;
}

bool CapturePropagation::optimizePartialApply(PartialApplyInst *PAI) {
   PILFunction *SubstF = PAI->getReferencedFunctionOrNull();
   if (!SubstF)
      return false;
   if (SubstF->isExternalDeclaration())
      return false;

   if (PAI->hasSubstitutions() && PAI->getSubstitutionMap().hasArchetypes()) {
      LLVM_DEBUG(llvm::dbgs()
                    << "CapturePropagation: cannot handle partial specialization "
                       "of partial_apply:\n";
                    PAI->dumpInContext());
      return false;
   }


   // First possibility: Is it a partial_apply where all partially applied
   // arguments are dead?
   std::pair<PILFunction *, PILFunction *> GenericSpecialized;
   PILOptFunctionBuilder FuncBuilder(*this);
   if (auto *NewFunc = getSpecializedWithDeadParams(FuncBuilder,
                                                    PAI, SubstF, PAI->getNumArguments(), GenericSpecialized)) {
      rewritePartialApply(PAI, NewFunc);
      if (GenericSpecialized.first) {
         // Notify the pass manager about the new function.
         addFunctionToPassManagerWorklist(GenericSpecialized.first,
                                          GenericSpecialized.second);
      }
      return true;
   }

   // Second possibility: Are all partially applied arguments constant?
   for (auto Arg : PAI->getArguments()) {
      if (!isConstant(Arg))
         return false;
   }
   if (!isProfitable(SubstF))
      return false;

   LLVM_DEBUG(llvm::dbgs() << "Specializing closure for constant arguments:\n"
                           << "  " << SubstF->getName() << "\n"
                           << *PAI);
   ++NumCapturesPropagated;
   PILFunction *NewF = specializeConstClosure(PAI, SubstF);
   rewritePartialApply(PAI, NewF);

   addFunctionToPassManagerWorklist(NewF, SubstF);
   return true;
}

void CapturePropagation::run() {
   DominanceAnalysis *DA = PM->getAnalysis<DominanceAnalysis>();
   auto *F = getFunction();
   bool HasChanged = false;

   // Don't optimize functions that are marked with the opt.never attribute.
   if (!F->shouldOptimize())
      return;

   // Cache cold blocks per function.
   ColdBlockInfo ColdBlocks(DA);
   for (auto &BB : *F) {
      if (ColdBlocks.isCold(&BB))
         continue;

      auto I = BB.begin();
      while (I != BB.end()) {
         PILInstruction *Inst = &*I;
         ++I;
         if (auto *PAI = dyn_cast<PartialApplyInst>(Inst))
            HasChanged |= optimizePartialApply(PAI);
      }
   }
   if (HasChanged) {
      invalidateAnalysis(PILAnalysis::InvalidationKind::Everything);
   }
}

PILTransform *polar::createCapturePropagation() {
   return new CapturePropagation();
}
